import torch
import json
import random
class DatasetIterater(object):
    def __init__(self, batches, batch_size, device, Trainflag=False):
        self.batch_size = batch_size
        self.batches = batches
        self.n_batches = len(batches) // batch_size
        self.residue = False  # 记录batch数量是否为整数
        if len(batches) % self.n_batches != 0:
            self.residue = True
        self.index = 0
        self.device = device
        self.train_flag = Trainflag

    def _to_batch(self, datas):
        doc = [str(_[0]) for _ in datas]
        query = [str(_[1]) for _ in datas]
        label = torch.LongTensor([_[2] for _ in datas]).to(self.device)

        # pad前的长度(超过pad_size的设为pad_size)
        return doc, query, label

    def __next__(self):
        if self.residue and self.index == self.n_batches:
            batches = self.batches[self.index * self.batch_size: len(self.batches)]
            self.index += 1
            batches = self._to_batch(batches)
            return batches

        elif self.index >= self.n_batches:
            self.index = 0
            if self.train_flag:
                random.shuffle(self.batches)
            raise StopIteration
        else:
            
            batches = self.batches[self.index * self.batch_size: (self.index + 1) * self.batch_size]
            self.index += 1
            batches = self._to_batch(batches)
            return batches

    def __iter__(self):
        return self

    def __len__(self):
        if self.residue:
            return self.n_batches + 1
        else:
            return self.n_batches


def build_iterator(dataset, batch_size, device,trainflag=False):
    iter = DatasetIterater(dataset, batch_size, device,trainflag)
    return iter

def build_data(filep):
    with open(filep,'r',encoding='utf-8') as fr:
        datas = []
        for line in fr:
            con = json.loads(line)
            if 'neg_responses' in con:
                datas.append([con["pos_response"],con["question"],1])
                datas.append([con["neg_responses"][0],con["question"],0])
            else:
                datas.append([con["pos_response"],con["question"],1])
                datas.append([con["neg_response"],con["question"],0])
        return datas
#     random.shuffle(datas)
